import pandas as pd
import matplotlib.pyplot as plt
import os

def plot_precision_analysis(base_dir, target_dir):
    """
    Plot relative_error and avg_relative_error vs proc_dir for different cn_dir groups
    
    Args:
        base_dir (str): Path to the avg_precision_summary.csv file
        target_dir (str): Directory to save the generated plots
    """
    # Create target directory if it doesn't exist
    os.makedirs(target_dir, exist_ok=True)
    # Read the CSV file
    df = pd.read_csv(base_dir)
    
    # Get unique cn_dir values
    cn_dirs = df['cn_dir'].unique()
    
    # Create figure with subplots
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # Define colors for different cn_dir groups
    colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown']
    
    # Plot relative_error
    for i, cn_dir in enumerate(cn_dirs):
        cn_data = df[df['cn_dir'] == cn_dir]
        # Sort by proc_dir to ensure proper line connection
        cn_data = cn_data.sort_values('proc_dir')
        
        # Extract numeric values from proc_dir (remove 'proc' suffix)
        proc_values = [int(proc.replace('proc', '')) for proc in cn_data['proc_dir']]
        
        ax1.plot(proc_values, cn_data['relative_error'], 
                marker='o', label=f'{cn_dir}', color=colors[i % len(colors)], linewidth=2)
    
    ax1.set_xlabel('Process Count')
    ax1.set_ylabel('Relative Error (%)')
    ax1.set_title('Relative Error vs Process Count')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    ax1.set_xscale('log', base=2)  # Log scale for better visualization
    
    # Plot avg_relative_error
    for i, cn_dir in enumerate(cn_dirs):
        cn_data = df[df['cn_dir'] == cn_dir]
        # Sort by proc_dir to ensure proper line connection
        cn_data = cn_data.sort_values('proc_dir')
        
        # Extract numeric values from proc_dir (remove 'proc' suffix)
        proc_values = [int(proc.replace('proc', '')) for proc in cn_data['proc_dir']]
        
        ax2.plot(proc_values, cn_data['avg_relative_error'], 
                marker='s', label=f'{cn_dir}', color=colors[i % len(colors)], linewidth=2)
    
    ax2.set_xlabel('Process Count')
    ax2.set_ylabel('Average Relative Error (%)')
    ax2.set_title('Average Relative Error vs Process Count')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    ax2.set_xscale('log', base=2)  # Log scale for better visualization
    
    # Adjust layout
    plt.tight_layout()
    
    # Save the plot
    plt.savefig(os.path.join(target_dir, 'precision_comparison.png'), dpi=300, bbox_inches='tight')
    plt.show()
    
    # Also create a combined plot showing both metrics for each cn_dir
    fig2, axes = plt.subplots(len(cn_dirs), 1, figsize=(12, 4 * len(cn_dirs)))
    if len(cn_dirs) == 1:
        axes = [axes]
    
    for i, cn_dir in enumerate(cn_dirs):
        cn_data = df[df['cn_dir'] == cn_dir]
        cn_data = cn_data.sort_values('proc_dir')
        
        proc_values = [int(proc.replace('proc', '')) for proc in cn_data['proc_dir']]
        
        axes[i].plot(proc_values, cn_data['relative_error'], 
                    marker='o', label='Relative Error', color='blue', linewidth=2)
        axes[i].plot(proc_values, cn_data['avg_relative_error'], 
                    marker='s', label='Avg Relative Error', color='red', linewidth=2)
        
        axes[i].set_xlabel('Process Count')
        axes[i].set_ylabel('Error (%)')
        axes[i].set_title(f'Error Metrics for {cn_dir}')
        axes[i].legend()
        axes[i].grid(True, alpha=0.3)
        axes[i].set_xscale('log', base=2)
    
    plt.tight_layout()
    plt.savefig(os.path.join(target_dir, 'precision_individual_plots.png'), dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    # Define the base directory path to avg_precision_summary.csv
    base_dir = r"F:\PostGraduate\Point-to-Point-DATA\deal-data-code\C-lop-Prediction\analysis_for_diff_network_topology\2node\all_predict_precision\avg_precision_summary.csv"
    
    # Define the target directory for saving plots
    target_dir = r"F:\PostGraduate\Point-to-Point-DATA\deal-data-code\C-lop-Prediction\analysis_for_diff_network_topology\2node\all_predict_precision/plots"
    
    # Check if file exists
    if os.path.exists(base_dir):
        plot_precision_analysis(base_dir, target_dir)
        print(f"Plots saved to: {target_dir}")
    else:
        print(f"Error: File not found at {base_dir}")
        print("Please check the file path and try again.")
